package code.bolin.ufw;

import code.bolin.ufw.common.NetUtils;
import code.bolin.ufw.common.UfwShell;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.InetAddress;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

import static code.bolin.ufw.common.LogUtil.log;

/**
 * @author bolin
 * @create 2023-05
 * @desc
 */
public class Main {

    static final Map<String, String> DOMAIN_MAP = new HashMap<>();
    static final Queue<String> waitDeleteRule = new LinkedBlockingQueue<>();
    static final Queue<String> waitAddRule = new LinkedBlockingQueue<>();

    static ScheduledThreadPoolExecutor poolExecutor = new ScheduledThreadPoolExecutor(1);

    public static void main(String[] args) throws InterruptedException {
        if (args == null || args.length < 1) {
            System.out.println("输入域名参数半角逗号分割,示例 baidu.com,qq.com");
            System.exit(1);
            return;
        }
        if (args[1].equalsIgnoreCase("deleteOld")) {
            deleteOldList();
        }
        for (String domainName : args[0].split(",")) {
            DOMAIN_MAP.put(domainName, null);
        }

        // 启动定时任务，更新ip
        startScheduleJobs_parseDomainAndNotifyAll();
        // 处理修改ufw任务
        processChangeUfwTaskQueue();

        // 阻塞主线程
        while (true) {
            synchronized (Thread.currentThread()) {
                Thread.currentThread().wait();
            }
        }
    }

    private static void processChangeUfwTaskQueue() {
        while (true) {
            if (!waitAddRule.isEmpty()) {
                String addIp = null;
                while ((addIp = waitAddRule.poll()) != null) {
                    execShell(UfwShell.ALLOW_ALL_FROM, addIp);
                    execShell(UfwShell.ROUTE_ALLOW_ALL_FROM, addIp);
                }
            } else if (!waitDeleteRule.isEmpty()) {
                String delIp = null;
                while ((delIp = waitDeleteRule.poll()) != null) {
                    execShell(UfwShell.DELETE_ALLOW_ALL_FROM, delIp);
                    execShell(UfwShell.DELETE_ROUTE_ALLOW_ALL_FROM, delIp);
                }
            } else {
                try {
                    synchronized (DOMAIN_MAP) {
                        DOMAIN_MAP.wait();
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 每隔10分钟解析一次domainMap中的域名，不一致的时候丢到待更新队列中
     */
    private static void startScheduleJobs_parseDomainAndNotifyAll() {
        poolExecutor.scheduleWithFixedDelay(() -> {
            DOMAIN_MAP.entrySet().stream().forEach(e -> {
                String domainName = e.getKey();
                String oldIp = e.getValue();
                try {
                    InetAddress address = InetAddress.getByName(domainName);
                    String newIp = address.getHostAddress();
                    if (!newIp.equals(oldIp)) {
                        log(String.format("网站 %s 的IP地址：%s >> %s", domainName, oldIp, newIp));
                        DOMAIN_MAP.put(domainName, newIp);
                        if (newIp != null) {
                            waitAddRule.offer(newIp);
                        }
                        if (oldIp != null) {
                            waitDeleteRule.offer(oldIp);
                        }
                        // 唤醒其他线程
                        synchronized (DOMAIN_MAP) {
                            DOMAIN_MAP.notifyAll();
                        }
                    }
                } catch (Exception exception) {
                    exception.printStackTrace();
                    log("无法解析该域名");
                }
            });
        }, 0, 1, TimeUnit.MINUTES);
    }


    /**
     * 删除非内网的from ip to Anywhere
     */
    private static void deleteOldList() {
        List<String> list = execShell(UfwShell.STATUS);
        List<String> waitDeleteRule = list.stream()
                .map(line -> line.split("ALLOW IN"))
                .filter(strs -> {
                    if (strs.length == 2) {
                        if (strs[0].indexOf("Anywhere") != -1) {
                            // 过滤内网IP
                            return !NetUtils.isPrivate(strs[1].trim());
                        }
                    }
                    return false;
                })
                .map(strs -> strs[1].trim())
                .collect(Collectors.toList());
        List<String> waitDeleteRule2 = list.stream()
                .map(line -> line.split("ALLOW FWD"))
                .filter(strs -> {
                    if (strs.length == 2) {
                        if (strs[0].indexOf("Anywhere") != -1) {
                            // 过滤内网IP
                            return !NetUtils.isPrivate(strs[1].trim());
                        }
                    }
                    return false;
                })
                .map(strs -> strs[1].trim())
                .collect(Collectors.toList());
        // 删除之前的域名
        waitDeleteRule.stream().forEach(rule -> execShell(UfwShell.DELETE_ALLOW_ALL_FROM, rule));
        waitDeleteRule2.stream().forEach(rule -> execShell(UfwShell.DELETE_ROUTE_ALLOW_ALL_FROM, rule));
    }

    /**
     * 执行shell指令 返回结果
     *
     * @param sh
     * @param params
     * @return
     */
    public static List<String> execShell(UfwShell sh, String... params) {
        log(String.format(String.format("%s: %s", sh.getName(), sh.getShell()), params));
        List<String> lines = new ArrayList<>();
        try {
            Process process = Runtime.getRuntime().exec(String.format(sh.getShell(), params));
            BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
            String line;
            while ((line = reader.readLine()) != null) {
                log(line);
                lines.add(line);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return lines;
    }

}
